import numpy as np
from . import LocalSearchOperator, register_operator


class ForwardInsertionOperator(LocalSearchOperator):
    """前向插入(Forward Insertion)局部搜索算子"""

    def __init__(self):
        self.name = "forward_insertion"
        self.supported_problems = ["tsp"]

    def generate_neighborhood(self, solution, problem_instance):
        """生成前向插入邻域"""
        neighborhood = []
        n = len(solution)

        # 对每个城市，尝试将其插入到所有可能的前向位置
        for i in range(1, n):  # 从第1个位置开始（第0个位置是起点）
            for j in range(0, i):  # 插入到i之前的位置
                if j == i or j == i - 1:  # 跳过相邻位置（效果等同于不移动）
                    continue

                # 创建新解：将城市i插入到位置j
                new_solution = solution.copy()
                city_to_move = new_solution[i]

                # 移除城市i
                new_solution = np.concatenate([new_solution[:i], new_solution[i + 1:]])

                # 插入到位置j
                new_solution = np.concatenate([new_solution[:j], [city_to_move], new_solution[j:]])

                neighborhood.append(new_solution)

        return neighborhood

    def efficient_forward_insertion(self, solution, problem_instance, intensity=15):
        """高效的前向插入实现，直接计算成本变化"""
        n = len(solution)
        current_solution = solution.copy()
        current_cost = problem_instance.evaluate(current_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False

            # 尝试多个随机插入操作
            for _ in range(min(100, n * n)):  # 限制尝试次数
                # 随机选择一个城市和一个插入位置
                i = np.random.randint(1, n)  # 要移动的城市位置（不能是第一个）
                j = np.random.randint(0, i)  # 插入位置（在i之前）

                if j == i or j == i - 1:  # 跳过无效位置
                    continue

                # 计算插入前的相关边成本
                prev_i = (i - 1) % n
                next_i = (i + 1) % n
                prev_j = (j - 1) % n

                # 计算插入前的成本
                old_cost = (
                        problem_instance.distance_matrix[current_solution[prev_i]][current_solution[i]] +
                        problem_instance.distance_matrix[current_solution[i]][current_solution[next_i]] +
                        problem_instance.distance_matrix[current_solution[prev_j]][current_solution[j]]
                )

                # 执行插入操作
                city_to_move = current_solution[i]
                new_solution = np.concatenate([
                    current_solution[:i],
                    current_solution[i + 1:]
                ])
                new_solution = np.concatenate([
                    new_solution[:j],
                    [city_to_move],
                    new_solution[j:]
                ])

                # 计算插入后的相关边成本
                new_prev_i = (j - 1) % (n - 1)  # 注意：移除一个城市后长度变为n-1
                new_next_i = j % (n - 1)
                new_next_j = (j + 1) % (n - 1)

                # 计算插入后的成本
                new_cost = (
                        problem_instance.distance_matrix[new_solution[new_prev_i]][new_solution[j]] +
                        problem_instance.distance_matrix[new_solution[j]][new_solution[new_next_i]] +
                        problem_instance.distance_matrix[new_solution[new_prev_i]][new_solution[new_next_j]]
                )

                # 计算成本变化（需要调整，因为边数发生了变化）
                # 更准确的方法是直接评估整个解，但计算成本高
                # 这里使用近似方法
                cost_change = new_cost - old_cost

                # 如果成本降低，则接受插入
                if cost_change < 0:
                    current_solution = new_solution
                    current_cost = problem_instance.evaluate(current_solution)  # 直接评估新解
                    improved = True
                    break  # 找到一个改进就跳出内层循环

            iterations += 1

        improvement = problem_instance.evaluate(solution) - current_cost
        return current_solution, improvement

    def __call__(self, solution, problem_instance, intensity=15):
        """应用前向插入局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        n = len(solution)
        if n < 4:  # 前向插入需要至少4个城市
            return solution, 0

        # 使用高效实现
        return self.efficient_forward_insertion(solution, problem_instance, intensity)


# 注册算子
register_operator("forward_insertion", ForwardInsertionOperator())